{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretnote.formal.locations import (\n",
    "    SPUFieldType,\n",
    "    SPUProtocolKind,\n",
    "    SymbolicPYU,\n",
    "    SymbolicSPU,\n",
    "    SymbolicWorld,\n",
    ")\n",
    "from secretflow.device import SPUCompilerNumReturnsPolicy\n",
    "\n",
    "sym_world = SymbolicWorld(world=frozenset((\"alice\", \"bob\")))\n",
    "sym_alice = SymbolicPYU(\"alice\")\n",
    "sym_bob = SymbolicPYU(\"bob\")\n",
    "sym_spu = SymbolicSPU(\n",
    "    world=frozenset((\"alice\", \"bob\")),\n",
    "    protocol=SPUProtocolKind.SEMI2K,\n",
    "    field=SPUFieldType.FM128,\n",
    "    fxp_fraction_bits=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "from secretnote.formal.locations import SFConfigSimulation, PortBinding\n",
    "\n",
    "ray.shutdown()\n",
    "\n",
    "sym_world.reify(SFConfigSimulation())\n",
    "\n",
    "alice = sym_alice.reify()\n",
    "bob = sym_bob.reify()\n",
    "spu = sym_spu.reify(\n",
    "    alice=PortBinding(announced_as=\"127.0.0.1:32767\"),\n",
    "    bob=PortBinding(announced_as=\"127.0.0.1:32768\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretnote.instrumentation.sdk import create_profiler, setup_tracing\n",
    "\n",
    "setup_tracing()\n",
    "\n",
    "profiler = create_profiler()\n",
    "profiler.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.datasets import load_breast_cancer\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "\n",
    "def load_dataset_for_training(rand_key: int, party_id: int):\n",
    "    x, y = load_breast_cancer(return_X_y=True)\n",
    "    x = (x - np.min(x)) / (np.max(x) - np.min(x))\n",
    "    x_train, x_test, y_train, y_test = train_test_split(\n",
    "        x,\n",
    "        y,\n",
    "        test_size=0.2,\n",
    "        random_state=rand_key,\n",
    "    )\n",
    "    if party_id == 0:\n",
    "        return x_train[:, :15], None\n",
    "    else:\n",
    "        return x_train[:, 15:], y_train\n",
    "\n",
    "\n",
    "def load_dataset_for_testing(rand_key: int):\n",
    "    x, y = load_breast_cancer(return_X_y=True)\n",
    "    x = (x - np.min(x)) / (np.max(x) - np.min(x))\n",
    "    x_train, x_test, y_train, y_test = train_test_split(\n",
    "        x,\n",
    "        y,\n",
    "        test_size=0.2,\n",
    "        random_state=rand_key,\n",
    "    )\n",
    "    return x_test, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    return 1 / (1 + jax.numpy.exp(-x))\n",
    "\n",
    "\n",
    "def predict(W, b, inputs):\n",
    "    \"\"\"Outputs probability of a label being true.\"\"\"\n",
    "    return sigmoid(jax.numpy.dot(inputs, W) + b)\n",
    "\n",
    "\n",
    "def loss(W, b, inputs, targets):\n",
    "    \"\"\"Training loss is the negative log-likelihood of the training examples.\"\"\"\n",
    "    preds = predict(W, b, inputs)\n",
    "    label_probs = preds * targets + (1 - preds) * (1 - targets)\n",
    "    return -jax.numpy.mean(jax.numpy.log(label_probs))\n",
    "\n",
    "\n",
    "def train_step(W, b, x1, x2, y, learning_rate):\n",
    "    x = jax.numpy.concatenate([x1, x2], axis=1)\n",
    "    Wb_grad = grad(loss, (0, 1))(W, b, x, y)\n",
    "    W -= learning_rate * Wb_grad[0]\n",
    "    b -= learning_rate * Wb_grad[1]\n",
    "    return W, b\n",
    "\n",
    "\n",
    "def fit(W, b, x1, x2, y, epochs=1, learning_rate=1e-2):\n",
    "    for _ in range(epochs):\n",
    "        W, b = train_step(W, b, x1, x2, y, learning_rate=learning_rate)\n",
    "    return W, b\n",
    "\n",
    "\n",
    "def concatenate_samples(x1, x2):\n",
    "    return jax.numpy.concatenate([x1, x2], axis=1)\n",
    "\n",
    "\n",
    "def grad(weights, bias, x, y, *, learning_rate):\n",
    "    gradients = jax.grad(loss, (0, 1))(weights, bias, x, y)\n",
    "    weights -= learning_rate * gradients[0]\n",
    "    bias -= learning_rate * gradients[1]\n",
    "    return weights, bias\n",
    "\n",
    "\n",
    "def validate_model(weights, bias, X_test, y_test):\n",
    "    y_pred = predict(weights, bias, X_test)\n",
    "    return roc_auc_score(y_test, y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1, _ = alice(load_dataset_for_training, num_returns=2)(rand_key=42, party_id=0)\n",
    "x2, y = bob(load_dataset_for_training, num_returns=2)(rand_key=42, party_id=1)\n",
    "\n",
    "x1 = x1.to(spu)\n",
    "x2 = x2.to(spu)\n",
    "y = y.to(spu)\n",
    "\n",
    "weights = jax.numpy.zeros((30,))\n",
    "bias = 0.0\n",
    "\n",
    "weights = alice(lambda x: x)(weights)\n",
    "weights = weights.to(spu)\n",
    "\n",
    "bias = alice(lambda x: x)(bias)\n",
    "bias = bias.to(spu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 10\n",
    "\n",
    "x = spu(concatenate_samples)(x1, x2)\n",
    "\n",
    "for _ in range(10):\n",
    "    weights, bias = spu(\n",
    "        grad,\n",
    "        num_returns_policy=SPUCompilerNumReturnsPolicy.FROM_USER,\n",
    "        user_specified_num_returns=2,\n",
    "    )(\n",
    "        weights,\n",
    "        bias,\n",
    "        x,\n",
    "        y,\n",
    "        learning_rate=1e-2,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretflow import reveal\n",
    "\n",
    "weights = weights.to(bob)\n",
    "bias = bias.to(bob)\n",
    "x_test, y_test = bob(load_dataset_for_testing, num_returns=2)(rand_key=42)\n",
    "auc = reveal(bob(validate_model)(weights, bias, x_test, y_test))\n",
    "\n",
    "print(f\"{weights=}\")\n",
    "print(f\"{bias=}\")\n",
    "print(f\"{auc=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from secretnote.display import visualize_run\n",
    "\n",
    "profiler.stop()\n",
    "visualize_run(profiler)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
